#!/usr/bin/env python

from optparse import OptionParser
from argparser import ArgParser
from ProxyGrabber import ProxyGrabber
import sys, threading

class MassProxyGrabber(object):
    def __init__(self, args):
        # Setting all arguments
        sites = args['sites']
        proxies = args['proxies']
        self.whois = args['whois']
        self.quiet = args['quiet']
        self.output = args['output']
        self.sthreads = args['sthreads']
        self.pthreads = args['pthreads']
        self.test_types = args['test_types']

        # Setting sites and proxies
        self.sites = self.fileorlistAssign(sites)
        self.proxies = self.fileorlistAssign(proxies)

        # Getting amount of proxies and sites
        self.sites_amount = self.dummyRun(self.sites)
        self.proxies_amount = self.dummyRun(self.proxies)

        # Thread variables for sites
        self.sites_threads = []
        self.sites_madethreads = 0
        self.sites_errors = []
        # Thread variables for proxies
        self.proxies_threads = []
        self.proxies_madethreads = 0
        self.proxies_errors = []
       
    def fileorlistAssign(self, var):
        if var != None:
            if var[0] == 'file':
                return self.openFile(var[1]) # File
            else:
                return var[1].split(',') # List
        else:
            return var
    def openFile(self, file_to_open):
        try:
            return open(file_to_open, 'r')
        except IOError:
            print '\n[-] Error: could not open %s to read from' % (file_to_open)
            sys.exit(1)
    
    def dummyRun(self, data):
        amount = 0
        if data != None:
            if type(data) == list:
                amount = len(data)
            else:
                for i in data.xreadlines():
                    amount += 1
                data.seek(0) # Returning to beginning of file
            return amount
        else:
            return None
    
    def runThreads(self):
        def error_handler(errors):
            if len(errors) == 0:
                return
            else:
                for i in errors:
                    print '\n[-] Error: %s' % (errors[i])
                else:
                    while True:
                        try:
                            tocont = raw_input('\nDo you want to continue anyway?: ')
                        except KeyboardInterrupt:
                            print '\n\n[*] Aborted'
                            sys.exit(1)
                        if tocont in ('y', 'yes'):
                            return
                        elif tocont in ('n', 'no'):
                            sys.exit(1)
                        else:
                            continue
        if self.sites_madethreads == self.sthreads:
            for i in xrange(self.sites_madethreads):
                self.sites_threads[i].start()
            for i in xrange(self.sites_madethreads):
                self.sites_threads[i].join() # Waiting for completion
                error_handler(self.sites_threads[i].errors) # Checking for errors in thread
            # Clearing variables
            self.sites_threads = []
            self.sites_madethreads = 0
            self.sites_errors = []

    def sitesThreads(self, website, number):
        site = website.strip()
        t = ProxyGrabber(site, number)
        self.sites_threads.append(t)
        self.sites_madethreads += 1
        # Starting threads ("hardcore" threads number managing xD )
        if self.sites_amount <= self.sthreads:
            self.sthreads = self.sites_amount
            self.runThreads()
        if self.sites_amount > self.sthreads:
            self.runThreads()
            self.s_threads_remaining = self.sites_amount - self.sthreads
            if self.s_threads_remaining <= self.sthreads:
                self.sthreads = self.s_threads_remaining
            
    def main(self):
        if (self.proxies == None) and (self.sites == None): # No Proxies - No Sites
            print '\n[*] Nothing to do - no sites or proxies input: exiting'
            sys.exit(1)

        elif (self.sites != None) and (self.proxies == None): # Only Sites
            if type(self.sites) == list:
                for i in self.sites:
                    self.sitesThreads(i, self.sites.index(i))
            else:
                n = 0
                for i in self.sites.xreadlines():
                    self.sitesThreads(i, n)
                    n += 1

        elif (self.proxies != None) and (self.sites == None): # Only Proxies
            pass
        else:
            pass
        print '\n[*] Done!'

if __name__=='__main__':
    # Try to use psyco to maximize speed, ingnore if it's not found
    try:
        import psyco
        psyco.full()
    except ImportError:
        pass
        
    # Getting and checking arguments for errors
    ap = ArgParser()
    toexit = False
    for i in ap.errors:
        if ap.errors[i][0] == True:
            if toexit == False:
                toexit = True
            print ap.errors[i][1]
    if toexit == True:
        sys.exit(1)
        
    # Initializing and passing arguments to main class
    mpg = MassProxyGrabber(ap.getArguments())
    # Starting main function
    mpg.main()
    
    # Testing
    for i in (mpg.sites, mpg.proxies, mpg.whois, mpg.quiet, mpg.output, mpg.sthreads, mpg.pthreads, mpg.test_types):
        print '%s\n' % (i)
